"""
Copyright 2020 Huawei Technologies Co., Ltd

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import json
import argparse
import traceback
from common import set_env
from common import atc_model

# number of code per time in distance_compute_flat
code_num = 16384


def arg_parse():
    """
    Parse arguements to the operator model
    """

    parser = argparse.ArgumentParser(
        description='generate aicore operator model')

    parser.add_argument("-d",
                        dest='dim',
                        default=512,
                        type=int,
                        help="Feature dimension")

    return parser.parse_args()


def get_dist_compute_flat_json(queries_num, dim, code_num, file_path):
    # write dist_compute_flat json
    dist_compute_obj = [{
        "op":
            "DistanceComputeFlat",
        "input_desc": [{
            "format": "ND",
            "shape": [queries_num, dim],
            "type": "float16"
        }, {
            "format": "ND",
            "shape": [code_num // 16, dim // 16, 16, 16],
            "type": "float16"
        }, {
            "format": "ND",
            "shape": [code_num],
            "type": "float16"
        }],
        "output_desc": [{
            "format": "ND",
            "shape": [queries_num, code_num],
            "type": "float16"
        }, {
            "format": "ND",
            "shape": [32],
            "type": "uint16"
        }]
    }]

    isObj = json.dumps(dist_compute_obj)
    fileObject = open(file_path, 'w')
    fileObject.write(isObj)
    fileObject.close()


def generate_by_atc():
    # generate directory of model
    dirs = './op_models/'
    if not os.path.exists(dirs):
        os.makedirs(dirs)

    atc_model('distance_compute_flat_op1')
    atc_model('distance_compute_flat_op2')
    atc_model('distance_compute_flat_op4')
    atc_model('distance_compute_flat_op6')
    atc_model('distance_compute_flat_op8')
    atc_model('distance_compute_flat_op12')
    atc_model('distance_compute_flat_op16')
    atc_model('distance_compute_flat_op18')
    atc_model('distance_compute_flat_op24')
    atc_model('distance_compute_flat_op30')
    atc_model('distance_compute_flat_op32')
    atc_model('distance_compute_flat_op36')
    atc_model('distance_compute_flat_op42')
    atc_model('distance_compute_flat_op48')
    atc_model('distance_compute_flat_op64')
    atc_model('distance_compute_flat_op96')


if __name__ == '__main__':
    args = arg_parse()
    set_env()

    dim = args.dim

    # generate directory of config
    dirs = './config/'
    if not os.path.exists(dirs):
        os.makedirs(dirs)

    get_dist_compute_flat_json(1, dim, code_num,
                               './config/distance_compute_flat_op1.json')
    get_dist_compute_flat_json(2, dim, code_num,
                               './config/distance_compute_flat_op2.json')
    get_dist_compute_flat_json(4, dim, code_num,
                               './config/distance_compute_flat_op4.json')
    get_dist_compute_flat_json(6, dim, code_num,
                               './config/distance_compute_flat_op6.json')
    get_dist_compute_flat_json(8, dim, code_num,
                               './config/distance_compute_flat_op8.json')
    get_dist_compute_flat_json(12, dim, code_num,
                               './config/distance_compute_flat_op12.json')
    get_dist_compute_flat_json(16, dim, code_num,
                               './config/distance_compute_flat_op16.json')
    get_dist_compute_flat_json(18, dim, code_num,
                               './config/distance_compute_flat_op18.json')
    get_dist_compute_flat_json(24, dim, code_num,
                               './config/distance_compute_flat_op24.json')
    get_dist_compute_flat_json(30, dim, code_num,
                               './config/distance_compute_flat_op30.json')
    get_dist_compute_flat_json(32, dim, code_num,
                               './config/distance_compute_flat_op32.json')
    get_dist_compute_flat_json(36, dim, code_num,
                               './config/distance_compute_flat_op36.json')
    get_dist_compute_flat_json(42, dim, code_num,
                               './config/distance_compute_flat_op42.json')
    get_dist_compute_flat_json(48, dim, code_num,
                               './config/distance_compute_flat_op48.json')
    get_dist_compute_flat_json(64, dim, code_num,
                               './config/distance_compute_flat_op64.json')
    get_dist_compute_flat_json(96, dim, code_num,
                               './config/distance_compute_flat_op96.json')

    try:
        generate_by_atc()
    except OSError:
        traceback.print_exc()
        print("make sure set right ASCEND_HOME environment variable")
    else:
        traceback.print_exc()
